import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch
from PIL import Image
import os
import numpy as np


def draw_arrow_and_save_high_quality(image_path, start, end, color, output_folder):
    """
    Draws an arrow on an image using matplotlib and saves it with high quality.

    :param image_path: Path to the input image.
    :param start: Start coordinate of the arrow (x, y).
    :param end: End coordinate of the arrow (x, y).
    :param color: Color of the arrow.
    :param output_folder: Directory to save the modified image.
    """
    # Ensure the output directory exists
    os.makedirs(output_folder, exist_ok=True)

    # Open the image using PIL and convert to RGB if it's not
    image = Image.open(image_path)
    if image.mode != 'RGB':
        image = image.convert('RGB')
    img_array = np.array(image)

    # Get image dimensions
    height, width = img_array.shape[:2]

    # Create a matplotlib figure with the same size as the original image
    # Use a high DPI for high-quality output
    dpi = 100
    figsize = width / dpi, height / dpi
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    plt.axis('off')  # No axes for clean image output
    fig.subplots_adjust(left=0, right=1, bottom=0, top=1)

    # Display the image
    ax.imshow(img_array)
    ax.set_xlim([0, width])
    ax.set_ylim([height, 0])

    # Draw an arrow
    ax.annotate("", xy=end, xytext=start,
                arrowprops=dict(arrowstyle="->", color=color, lw=2))

    # Save the figure with high quality
    output_image_path = os.path.join(output_folder, f"{os.path.splitext(os.path.basename(image_path))[0]}_arrow.png")
    fig.savefig(output_image_path, dpi=dpi, format='png', bbox_inches='tight', pad_inches=0,
                transparent=True)

    # Close the figure to free memory
    plt.close(fig)

    return output_image_path

# Example usage
# Example usage:
image_paths = [
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\bUnwarp_wrap_crop1.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\elastic_wrap_crop1.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\emreg_wrap_crop1.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\ours_wrap_crop1.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\ssemnet_wrap_crop1.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\wrapgt_crop1.png',
    # 'image2.png',
    # Add your other image paths here
]
start_point = (70, 30)  # Starting point of the arrow
end_point = (50, 2)  # Ending point of the arrow
output_folder = 'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\指定区域\\箭头'
# Call the function for each image
for image_path in image_paths:
    draw_arrow_and_save_high_quality(image_path, start_point, end_point,"red" ,output_folder)
